import os
import faiss
import pickle
import numpy as np

class KnowledgeBase:
    _instance = None

    def __new__(cls, index_path, chunks_path, embeddings):
        if cls._instance is None:
            cls._instance = super(KnowledgeBase, cls).__new__(cls)
            cls._instance._initialized = False
        return cls._instance

    def __init__(self, index_path, chunks_path, embeddings):
        if self._initialized:
            return
        self.index_path = index_path
        self.chunks_path = chunks_path
        self.embeddings = embeddings
        self.index = None
        self.chunks = None
        self._load_local()
        self._initialized = True

    def _load_local(self):
        if os.path.exists(self.index_path) and os.path.exists(self.chunks_path):
            self.index = faiss.read_index(self.index_path)
            with open(self.chunks_path, "rb") as f:
                self.chunks = pickle.load(f)
        else:
            self.index = None
            self.chunks = None

    def save_local(self, chunks):
        print("正在生成文本向量...")
        vectors = self.embeddings.embed_documents(chunks)
        vector_dim = len(vectors[0])
        index = faiss.IndexFlatL2(vector_dim)
        index.add(np.array(vectors, dtype=np.float32))
        faiss.write_index(index, self.index_path)
        with open(self.chunks_path, "wb") as f:
            pickle.dump(chunks, f)
        self.index = index
        self.chunks = chunks
        print(f"FAISS索引創建完成，共存儲 {index.ntotal} 個向量（維度：{vector_dim}）")
        print(f"向量索引和文本塊已保存到本地（{self.index_path} 和 {self.chunks_path}）")

    def search(self, query, top_k=2):
        if self.index is None or self.chunks is None:
            print("未加載知識庫，請先執行 save_local 或確認本地文件存在")
            return []
        query_vector = self.embeddings.embed_query(query)
        query_vector_np = np.array([query_vector], dtype=np.float32)
        distances, indices = self.index.search(query_vector_np, top_k)
        results = []
        for i, idx in enumerate(indices[0]):
            results.append({
                "text": self.chunks[idx],
                "distance": distances[0][i]
            })
        return results 